#include <vector>
#include <unordered_map>
#include <unordered_set>
#include <cstring>
#include <libaio.h>
#include <iostream>
#include <algorithm>

#define MAX_BATCH_SIZE 128

class IOUringManager {
private:
    int fd_;
    int numRequests_;
    io_context_t ctx_;
    std::unordered_map<struct iocb*, int> cb_to_id_;
    std::unordered_set<struct iocb*> pending_iocbs_;

public:
    IOUringManager(int fd, int request_count)
        : fd_(fd), numRequests_(request_count) {
        memset(&ctx_, 0, sizeof(io_context_t));
        int ret = io_setup(numRequests_, &ctx_);
        if (ret < 0) {
            std::cerr << "Failed to initialize libaio. Error code: " << ret << std::endl;
            exit(EXIT_FAILURE);
        }
    }

    ~IOUringManager() {
        // Cleanup all pending I/O control blocks
        for (auto cb : pending_iocbs_) {
            delete cb;
        }
        pending_iocbs_.clear();
        cb_to_id_.clear();
        io_destroy(ctx_);
    }

    int get_fd() {
        return fd_;
    }

    int submitRequests(std::vector<void*>& input_bufAddrs,
                      const std::vector<int>& offsets,
                      const std::vector<int>& sizes,
                      const std::vector<int>& ids,
                      int input_num_request = 0) {
        size_t num_requests = input_num_request == 0 ? input_bufAddrs.size() : input_num_request;
        if (num_requests == 0) return 0;

        // Validate input vectors
        if (input_bufAddrs.size() < num_requests || offsets.size() < num_requests ||
            sizes.size() < num_requests || ids.size() < num_requests) {
            std::cerr << "Input vectors too small for requested operations" << std::endl;
            return 0;
        }

        std::vector<struct iocb*> iocb_ptrs;
        iocb_ptrs.reserve(num_requests);

        try {
            for (size_t i = 0; i < num_requests; ++i) {
                struct iocb* cb = new struct iocb;
                io_prep_pread(cb, fd_, input_bufAddrs[i], sizes[i], offsets[i]);
                cb->data = cb; // Store self pointer for cleanup
                cb_to_id_[cb] = ids[i];
                pending_iocbs_.insert(cb);
                iocb_ptrs.push_back(cb);
            }
        } catch (...) {
            // Cleanup if memory allocation fails
            for (auto cb : iocb_ptrs) {
                pending_iocbs_.erase(cb);
                cb_to_id_.erase(cb);
                delete cb;
            }
            return 0;
        }

        int ret = io_submit(ctx_, num_requests, iocb_ptrs.data());
        if (ret < 0) {
            // Full cleanup on submission failure
            for (auto cb : iocb_ptrs) {
                pending_iocbs_.erase(cb);
                cb_to_id_.erase(cb);
                delete cb;
            }
            return 0;
        } else if (ret < static_cast<int>(num_requests)) {
            // Partial success cleanup
            int submitted = ret;
            for (size_t i = submitted; i < num_requests; ++i) {
                struct iocb* cb = iocb_ptrs[i];
                pending_iocbs_.erase(cb);
                cb_to_id_.erase(cb);
                delete cb;
            }
            return submitted;
        }
        return ret;
    }

    int getCompletedResult(std::vector<int>& completedIds, int min_nr = 0) {
        struct io_event events[MAX_BATCH_SIZE];
        int ret = io_getevents(ctx_, min_nr, MAX_BATCH_SIZE, events, nullptr);
        if (ret < 0) {
            std::cerr << "Error retrieving events: " << ret << std::endl;
            return 0;
        }

        for (int i = 0; i < ret; ++i) {
            struct iocb* cb = static_cast<struct iocb*>(events[i].data);
            auto it = cb_to_id_.find(cb);
            if (it != cb_to_id_.end()) {
                completedIds.push_back(it->second);
                pending_iocbs_.erase(cb);
                cb_to_id_.erase(cb);
                delete cb;
            } else {
                std::cerr << "Orphaned I/O completion detected" << std::endl;
            }
        }
        return ret;
    }
};